"""
模块名：hxpITS
功能：软件后端顶层类
作者：hxp@海小皮
版本： V1.1
依赖：
    hxpITT
    hxpTTS
    DataStruct
    threading
    time
    SequenceMatcher
    soundfile
    sounddevice
"""
from typing import Dict
from .hxpITT import hxpITT, SCREEN_RESOLUTION_PARAM_DICT
from .DataStruct import GameType, ScreenType, Dialog, EmotionType
from .hxpTTS import hxpTTS, InferenceParam, Role
from threading import Lock, Thread, local
import time
from difflib import SequenceMatcher
import soundfile as sf
import sounddevice as sd
from collections import deque
import PlotDub.TextProcess as tp
from concurrent.futures import ThreadPoolExecutor, as_completed
from queue import Queue
import copy


class hxpITS:
    """
    Image To Speach
    功能：提供前端调用的 API，实现前端界面功能，主要实现剧情截图文字识别与配音功能
    成员：
        __ITT hxpITT 对象，用于捕获图像和图像转文本对话结构 Dialog
        __TTS hxpTTS 对象，用于 Dialog 对话文本转语音
        __ITSThreadSW ITS 任务线程开关
        __ITSThread ITS 任务线程
        __tryListenDict 试听字典 key = game+Role.name value = Role.emotionsDict
    """
    def __init__(self, games_roles_dict: Dict[str, Role], inference_param: InferenceParam):
        """
        功能：初始化 ITS
        输入：
            games_roles_dict 游戏角色模型及情感参数字典 key = Role.game.value+Role.name value = Role 对象
            inference_param 模型推理参数 InferenceParam 结构
        """
        self.__ITT = hxpITT()  # 默认为 1080p
        self.__TTS = hxpTTS(games_roles_dict, inference_param)  # hxpTTS 对象 用于文本转语音
        self.__ITSThreadSW = False  # ITS 任务线程开关
        self.__readQueue = deque()  # 待合成语音的分句队列
        self.__audioQueue = deque()  # 待播放音频队列
        self.__readDict: dict[str, bool] = {}  # 已读字典
        self.__ITSThread = None  # ITS 任务线程
        # 初始化 tryListenDict
        self.__tryListenDict = {}
        for game_role in games_roles_dict:
            # 读取 game_roles_dict 对应 game_role 的 emotionDict 存入 __tryListenDict
            self.__tryListenDict[game_role] = games_roles_dict[game_role].emotionsDict

        # 共享数据互斥锁
        self.__readDictLock = Lock()  # 保护 __readDict 的锁
        self.__readQueueLock = Lock()  # 保护 __readQueue 的锁
        self.__audioQueueLock = Lock()  # 保护 __audioQueue 的锁

        # TTS 线程池 (暂不可用)
        self.__localStorage = local()

    def TryListen(self, game: GameType, role_name, emo: EmotionType):
        """
        功能：试听角色参考音频
        输入：
            game GameType 枚举类型的游戏名
            role_name 角色名
            emotion EmotionType 枚举类型的情感
        输出：
            执行结果 True 找到对应音频并播放 False 未找到对应音频
        """
        #  查找 __tryListenDict 字典 获取 game+role 的 emotionDict
        emotion_dict = self.__tryListenDict.get(game.value + role_name)
        if emotion_dict:  # 存在对应角色
            audio_path = emotion_dict[emo].refWav  # 获取音频文件路径
            # 播放音频
            data, samplerate = sf.read(audio_path)
            sd.play(data, samplerate)
            return True
        else:  # 不存在对应角色
            print(f"【{self.__class__.__name__}】【{self.__tryListenDict.__name__}】未找到【{game.value}】【role】【{emo.value}】的音频！")
            return False

    def TryMerge(self, game: GameType, role_name, emo: EmotionType, text: str):
        """
        功能：调用 hxpTTS 的 TTS 进行一次直接推理
        输入：
            game GameType 枚举类型游戏名
            role_name 角色名
            emo EmotionType 枚举类型情感
            text 目标合成文字
        输出：合成结果
        """
        # 调用 TTS 进行推理
        audio = self.__TTS.TTS(game, role_name, emo, text)
        if audio:
            print(f"【{self.__class__.__name__}】【{self.TryMerge.__name__}】【{game.value}】【{role_name}】【{emo.value}】【{text}】合成成功！")
            return audio
        else:
            print(f"【{self.__class__.__name__}】【{self.TryMerge.__name__}】【{game.value}】【{role_name}】【{emo.value}】合成失败！")
            return None

    def AddGameRole(self, role: Role):
        """
        功能：更新 __TTS 对象的 gameRolesDict
        输入：
            role Role 对象
        """
        # 更新 TTS 对象
        self.__TTS.AddGameRole(role)
        # 更新试听字典
        self.__tryListenDict[role.game.value + role.name] = role.emotionsDict

    def SaveScreenResolution(self, resolution: ScreenType):
        """
        功能：将 resolution 参数映射到屏幕参数字典，传递给 ITT 模块完成分辨率的参数修改
        输入：
            resolution 屏幕分辨率参数字符串 eg “1920*1080”
        """
        self.__ITT.ResetShotParam(SCREEN_RESOLUTION_PARAM_DICT[resolution])
        print(f"【{self.__class__.__name__}】【{self.SaveScreenResolution.__name__}】新的屏幕分辨率参【{resolution.value}】已设置。")

    def SaveInferenceParam(self, inference_param: InferenceParam):
        """
        功能：推理参数保存，将推理函数保存到 TTS 内部
        输入：
            inference_param InferenceParam 对象
        """
        self.__TTS.SetParam(inference_param)
        print(f"【{self.__class__.__name__}】【{self.SaveInferenceParam.__name__}】新的模型推理参数已设置。")

    def __ITT_thread(self):
        """
        图像转文本线程
        共享数据：
            self.__readDict R & W
            self.__readQueue R & W
            self.__audioQueue R & W
        """
        detect_rate = 0.5  # 识别频率 单位 s
        last_dialog = Dialog("", "", EmotionType.CALM, [""])
        while self.__ITSThreadSW:
            # ITT 获取屏幕截图并转对话文本
            frame = self.__ITT.ScreenShot()
            dialog = self.__ITT.ImageToDialog(frame)

            # 识别有效性检测
            if dialog.roleName == "None":
                time.sleep(detect_rate)
                continue  # 识别无效，进入下一次检测

            # 判断是否是新的对话
            last_dialog_content = "".join(last_dialog.speachContent)
            dialog_content = "".join(dialog.speachContent)
            # print(f"{last_dialog_content} {dialog_content}")
            similarity = SequenceMatcher(None, dialog_content, last_dialog_content).ratio()
            last_dialog = dialog  # 更新 last
            if similarity < 0.1:  # 是新的对话
                with self.__readDictLock:
                    self.__readDict.clear()  # 清空阅读字典，准备新的对话检测
                with self.__readQueueLock:
                    self.__readQueue.clear()  # 清空阅读队列，准备新的对话阅读
                with self.__audioQueueLock:
                    self.__audioQueue.clear()  # 清空音频队列，准备新的音频播放
            else:  # 仍然是当前对话
                # 将新的分句添加到阅读字典中
                # print(f"【{self.__class__.__name__}】【ITT】【{str(dialog.speachContent)}】")
                for text in dialog.speachContent:
                    # 不是完整的分句直接跳过
                    if not tp.is_complete_sentence(text):
                        continue

                    # 遍历已读字典，检查重叠或重复内容
                    is_redundant = False
                    with self.__readDictLock:
                        read_dict = self.__readDict
                    for existing_text in read_dict.keys():
                        # 检查新句是否为已有句的子句，或已有句是否为新句的子句
                        if text == existing_text:
                            # print("text == existing_text")
                            pass
                        elif text in existing_text:
                            # print("text in existing_text")
                            is_redundant = True
                            break
                        elif existing_text in text:
                            # print("existing_text in text")
                            text = text.replace(existing_text, "")

                    if is_redundant:
                        # print(f"当前分句 {text} 已经被朗读过")
                        continue

                    # 若阅读字典中不存在该文本
                    if text not in read_dict and len(text) > 0:
                        with self.__readDictLock:
                            self.__readDict[text] = False  # 添加并标记为未读
                        with self.__readQueueLock:
                            self.__readQueue.append((dialog, text))  # 添加到阅读队列中
                        print(f"【{self.__class__.__name__}】【ITT】【{text}】添加到待合成队列")

            time.sleep(detect_rate)

    def __get_local_TTS(self):
        # 如果当前线程没有绑定的 TTS 实例，则创建一个
        if not hasattr(self.__localStorage, "TTS"):
            self.__localStorage.TTS = copy.deepcopy(self.__TTS)  # 创建一个独立的副本
        return self.__localStorage.TTS

    def __tts_task(self, game, index, dialog, read_text):
        """
        每个线程池任务调用的函数。
        """
        # 获取当前线程的独立 TTS 对象
        local_TTS = self.__get_local_TTS()
        # 调用独立的 __TTS 对象的 TTS 方法
        sampling_rate, audio_data = local_TTS.TTS(game, dialog.roleName, dialog.emotion, read_text)
        # 返回结果
        return index, read_text, sampling_rate, audio_data

    def __TTS_thread(self, game: GameType):
        """
        语音合成线程
        共享数据：
            self.__readDict W
            self.__readQueue R & W
            self.__audioQueue W
        """
        while self.__ITSThreadSW:
            # 从阅读队列中取出分句
            read_text = None
            with self.__readQueueLock:
                if len(self.__readQueue) > 0:
                    dialog, read_text = self.__readQueue.popleft()

            if read_text:
                # 调用 TTS 接口生成语音
                sampling_rate, audio_data = self.__TTS.TTS(game, dialog.roleName, dialog.emotion, read_text)
                # 将生成的语音数据放入音频队列
                with self.__audioQueueLock:
                    self.__audioQueue.append((read_text, sampling_rate, audio_data))
                # 更新已读状态
                with self.__readDictLock:
                    self.__readDict[read_text] = True
                print(f"【{self.__class__.__name__}】【TTS】【{read_text}】音频合成完成")
            else:
                time.sleep(0.1)  # 队列为空时稍作等待

        # 线程池方案 由于 TTS 类中的模型无法做到线程安全，因此内存开销极大，我的 16G 内存直接撑爆
        # max_workers = 3  # 最大并行任务数
        # task_queue = Queue()  # 临时存储处理结果的队列，用于排序
        #
        # while self.__ITSThreadSW:
        #     tasks = []
        #     with self.__readQueueLock:
        #         # 从队列中取最多 `max_workers` 个任务
        #         for i in range(max_workers):
        #             if len(self.__readQueue) > 0:
        #                 dialog, read_text = self.__readQueue.popleft()
        #                 tasks.append((i, dialog, read_text))
        #
        #     if tasks:
        #         # 使用线程池并行处理任务
        #         with ThreadPoolExecutor(max_workers=max_workers) as executor:
        #             futures = {
        #                 executor.submit(self.__tts_task, game, i, dialog, read_text): i
        #                 for i, dialog, read_text in tasks
        #             }
        #             for future in as_completed(futures):
        #                 index, read_text, sampling_rate, audio_data = future.result()
        #                 task_queue.put((index, read_text, sampling_rate, audio_data))
        #
        #         # 按照任务的提交顺序，将结果写入音频队列
        #         while not task_queue.empty():
        #             index, read_text, sampling_rate, audio_data = task_queue.get()
        #             with self.__audioQueueLock:
        #                 self.__audioQueue.append((read_text, sampling_rate, audio_data))
        #             # 更新已读状态
        #             with self.__readDictLock:
        #                 self.__readDict[read_text] = True
        #             print(f"【{self.__class__.__name__}】【TTS】【{read_text}】音频合成完成")
        #     else:
        #         time.sleep(0.1)  # 队列为空时稍作等待

    def __audio_play(self):
        """
        音频播放线程
        阻塞式播放，当前音频未播放完成时不会读取下一个音频
        共享数据：
            self.__audioQueue R & W
        """
        while self.__ITSThreadSW:
            can_play = False
            with self.__audioQueueLock:
                if len(self.__audioQueue) > 0:
                    # 从音频队列中取出语音数据
                    read_text, sampling_rate, audio_data = self.__audioQueue.popleft()
                    can_play = True  # 在这里放掉锁，防止音频播放阻塞其他线程

            if can_play:
                # 播放音频（阻塞式）
                sd.play(audio_data, sampling_rate)
                sd.wait()  # 等待音频播放完成
                print(f"【{self.__class__.__name__}】【audio_play】{read_text}音频朗读完成")
            else:
                time.sleep(0.1)  # 队列为空时稍作等待

    def __RunTask(self, game: GameType):
        """
        剧情配音任务线程
        输入：
            game GameType 枚举类型游戏名
        """
        # 启动 ITT 线程
        itt_thread = Thread(target=self.__ITT_thread, daemon=True)
        itt_thread.start()

        # 启动 TTS 线程
        tts_thread = Thread(target=self.__TTS_thread, args=(game,), daemon=True)
        tts_thread.start()

        # 启动音频播放线程
        audio_thread = Thread(target=self.__audio_play, daemon=True)
        audio_thread.start()

        # 等待线程执行
        itt_thread.join()
        tts_thread.join()
        audio_thread.join()

    def Run(self, game: GameType):
        """
        功能：启动剧情文本配音功能
        输入：
            game 要配音的游戏 GameType 枚举类型
        """
        if self.__ITSThread and self.__ITSThread.is_alive():
            print(f"【{self.__class__.__name__}】【{self.Run.__name__}】ITS 线程已经运行，无需重复启动！")
            return

        self.__ITSThreadSW = True
        self.__ITSThread = Thread(target=self.__RunTask, args=(game,))
        self.__ITSThread.start()
        print(f"【{self.__class__.__name__}】【{self.Run.__name__}】ITS 线程启动，【{game.value}】剧情配音开始。")

    def Stop(self):
        """
        功能：停止剧情文本配音功能
        """
        if not self.__ITSThread or not self.__ITSThread.is_alive():
            print(f"【{self.__class__.__name__}】【{self.Stop.__name__}】ITS 线程还未启动.")
            return

        self.__ITSThreadSW = False
        self.__ITSThread.join()
        print(f"【{self.__class__.__name__}】【{self.Stop.__name__}】ITS 线程已停止。")
